#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RTSP视频流接收器客户端
"""

import cv2
import numpy as np
import time
import argparse
import os
import threading
from PIL import Image, ImageDraw, ImageFont
from rtsp.rtsp_receiver import RTSPReceiver
import logging

logger = logging.getLogger('RTSPClient')
logger.setLevel(logging.DEBUG)

class RTSPClient:
    def __init__(self, config):
        # 加载中文字体
        try:
            # 尝试加载系统中的中文字体
            font_paths = [
                '/System/Library/Fonts/PingFang.ttc',  # macOS
                '/usr/share/fonts/truetype/wqy/wqy-microhei.ttc',  # Ubuntu
                'C:/Windows/Fonts/simhei.ttf',  # Windows
                'simhei.ttf'  # 当前目录
            ]
            
            self.font_path = None
            for path in font_paths:
                if os.path.exists(path):
                    self.font_path = path
                    break
            
            if self.font_path is None:
                print("警告：未找到中文字体文件，将使用默认字体，中文可能显示为乱码")
        except Exception as e:
            print(f"加载字体时出错: {str(e)}")
            self.font_path = None
        
        # 初始化成员变量
        self.config = config
        self.receiver = None
        self.stream_thread = None
        self.running = False
        self.processing_mode = 0  # 0: 原始, 1: 灰度, 2: 边缘检测, 3: 模糊
        self.is_recording = False
        self.video_writer = None
        self.current_frame = None
        self.frame_lock = threading.Lock()  # 用于线程同步
        
        # 初始化RTSP接收器
        # 创建输出目录
        self.output_dir = self.create_output_dir(config.get('output', 'output'))
        
        # 创建RTSP接收器
        self.receiver = RTSPReceiver(
            rtsp_url=config.get('url'),
            buffer_size=config.get('buffer', 10),
            reconnect_attempts=config.get('reconnect', 5),
            reconnect_delay=config.get('delay', 2.0)
        )
        
        # 初始化窗口
        self.window_name = "RTSP Stream Viewer"
        
        # 打印控制信息
        print("\n" + "="*50)
        print("RTSP视频流接收器")
        print("="*50)
        print("按键控制:")
        print("  Q: 退出")
        print("  S: 保存当前帧")
        print("  P: 暂停/恢复")
        print("  R: 开始/停止录制")
        print("  1: 原始视频")
        print("  2: 灰度效果")
        print("  3: 边缘检测")
        print("  4: 模糊效果")
        print("="*50 + "\n")
        # 在init方法中不自动启动流，而是在主程序中显式启动
        self.start_stream_thread()

    @staticmethod
    def parse_arguments():
        """解析命令行参数"""
        parser = argparse.ArgumentParser(description='RTSP视频流接收器')
        parser.add_argument('--url', type=str, required=True, help='RTSP流URL')
        parser.add_argument('--buffer', type=int, default=10, help='帧缓冲区大小')
        parser.add_argument('--reconnect', type=int, default=5, help='重连尝试次数')
        parser.add_argument('--delay', type=float, default=2.0, help='重连延迟(秒)')
        parser.add_argument('--output', type=str, default='output', help='输出目录')
        parser.add_argument('--record', action='store_true', help='是否录制视频')
        return parser.parse_args()

    @staticmethod
    def create_output_dir(output_dir):
        """创建输出目录"""
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"创建输出目录: {output_dir}")
        return output_dir

    def apply_edge_detection(self, frame):
        """应用边缘检测"""
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        edges = cv2.Canny(blurred, 50, 150)
        return cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)

    def apply_grayscale(self, frame):
        """应用灰度转换"""
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)

    def apply_blur(self, frame):
        """应用模糊效果"""
        return cv2.GaussianBlur(frame, (15, 15), 0)

    def resize_frame(self, frame, width=None, height=None):
        """调整帧大小"""
        if width is None and height is None:
            return frame
        
        h, w = frame.shape[:2]
        if width is None:
            aspect = height / float(h)
            dim = (int(w * aspect), height)
        elif height is None:
            aspect = width / float(w)
            dim = (width, int(h * aspect))
        else:
            dim = (width, height)
        
        return cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)

    def put_chinese_text(self, img, text, position, font_size=20, color=(0, 255, 0)):
        """使用PIL绘制中文文本"""
        # 如果是OpenCV图像格式(BGR)，转换为RGB
        if isinstance(img, np.ndarray):
            pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        else:
            pil_img = img
        
        # 创建绘图对象
        draw = ImageDraw.Draw(pil_img)
        
        # 加载字体
        try:
            if self.font_path:
                font = ImageFont.truetype(self.font_path, font_size)
            else:
                # 使用默认字体
                font = ImageFont.load_default()
        except Exception as e:
            print(f"加载字体失败: {str(e)}")
            font = ImageFont.load_default()
        
        # 绘制文本
        draw.text(position, text, font=font, fill=color)
        
        # 如果输入是OpenCV图像，转换回OpenCV格式
        if isinstance(img, np.ndarray):
            return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
        else:
            return pil_img

    def draw_info_on_frame(self, frame, info):
        """在帧上绘制信息"""
        h, w = frame.shape[:2]
        
        # 创建一个副本，避免修改原始帧
        display_frame = frame.copy()
        
        info_text = [
            f"分辨率: {w}x{h}",
            f"FPS: {info['fps']:.2f}",
            f"帧计数: {info['frame_count']}",
            f"状态: {'暂停中' if info['is_paused'] else '运行中'}",
            f"连接: {'已连接' if info['connection_status'] else '未连接'}"
        ]
        
        # 使用PIL绘制中文文本
        for i, text in enumerate(info_text):
            y_pos = 30 + i * 30
            display_frame = self.put_chinese_text(display_frame, text, (10, y_pos - 20), font_size=20)
        
        # 添加控制说明
        controls = [
            "按键控制:",
            "Q: 退出",
            "S: 保存当前帧",
            "P: 暂停/恢复",
            "R: 开始/停止录制",
            "1: 原始视频",
            "2: 灰度效果",
            "3: 边缘检测",
            "4: 模糊效果"
        ]
        
        for i, text in enumerate(controls):
            y_pos = 30 + i * 30
            display_frame = self.put_chinese_text(display_frame, text, (w - 200, y_pos - 20), font_size=18, color=(255, 255, 0))
        
        return display_frame

    def start_stream_thread(self, show_window=False):
        """在单独的线程中启动视频流循环"""
        if self.running:
            print("视频流已经在运行中")
            return False
        
        # 连接并开始接收
        if not self.receiver.connect():
            print("无法连接到RTSP流")
            return False
        
        self.receiver.start()
        
        # 创建窗口（如果需要显示）
        if show_window:
            # 在macOS上，使用WINDOW_NORMAL可能更可靠
            cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        
        # 设置运行标志
        self.running = True
        
        # 创建并启动线程
        # self.stream_loop(show_window)
        self.stream_thread = threading.Thread(target=self.stream_loop, args=(show_window,))
        self.stream_thread.daemon = True  # 设置为守护线程，主线程结束时自动结束
        self.stream_thread.start()
        print("视频流线程已启动")
        return True
    
    def stream_loop(self, show_window=True):
        """视频流处理的主循环"""
        try:
            
            while self.running:
                # 获取当前帧
                frame = self.receiver.get_current_frame()
                
                if frame is not None:
                    # 使用线程锁保护当前帧的更新
                    with self.frame_lock:
                        self.current_frame = frame.copy()
                    
                    if show_window:
                        # 应用选定的处理模式
                        processed_frame = frame.copy()
                        if self.processing_mode == 1:
                            processed_frame = self.apply_grayscale(processed_frame)
                        elif self.processing_mode == 2:
                            processed_frame = self.apply_edge_detection(processed_frame)
                        elif self.processing_mode == 3:
                            processed_frame = self.apply_blur(processed_frame)
                        
                        # 获取并显示帧信息
                        info = self.receiver.get_frame_info()
                        display_frame = self.draw_info_on_frame(processed_frame.copy(), info)
                        
                        # 确保帧不为空且尺寸合适
                        if display_frame is not None and display_frame.size > 0:
                            # 显示帧
                            cv2.imshow(self.window_name, display_frame)
                            
                            # 在macOS上，使用更长的等待时间可能有助于窗口显示
                            key = cv2.waitKey(30)
                            self.process_key_input(key, frame)
                        
                        # 如果正在录制，写入视频
                        if self.is_recording and self.video_writer is not None:
                            self.video_writer.write(processed_frame)
                else:
                    # 如果没有帧，短暂休眠以减少CPU使用
                    time.sleep(0.01)
                    
        except Exception as e:
            print(f"视频流线程发生错误: {str(e)}")
        finally:
            # 清理资源
            if self.is_recording and self.video_writer is not None:
                self.video_writer.release()
                self.video_writer = None
                self.is_recording = False
            
            if show_window:
                cv2.destroyAllWindows()
            
            self.receiver.disconnect()
            self.running = False
            print("视频流线程已停止")
    
    def process_key_input(self, key, frame):
        """处理键盘输入"""
        if key == ord('q'):
            print("退出程序")
            self.stop_stream()
        elif key == ord('s'):
            # 使用新的save_current_frame方法保存当前帧
            self.save_current_frame()
        elif key == ord('p'):
            # 暂停/恢复
            if self.receiver.is_paused():
                self.receiver.resume()
                print("恢复视频流")
            else:
                self.receiver.pause()
                print("暂停视频流")
        elif key == ord('r'):
            # 开始/停止录制
            if self.is_recording:
                self.is_recording = False
                if self.video_writer is not None:
                    self.video_writer.release()
                    self.video_writer = None
                print("停止录制")
            else:
                # 获取视频参数
                if frame is not None:
                    h, w = frame.shape[:2]
                    info = self.receiver.get_frame_info()
                    fps = info['fps'] if info['fps'] > 0 else 30
                    timestamp = time.strftime("%Y%m%d_%H%M%S")
                    video_path = os.path.join(self.output_dir, f"video_{timestamp}.mp4")
                    
                    # 创建视频写入器
                    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                    self.video_writer = cv2.VideoWriter(video_path, fourcc, fps, (w, h))
                    
                    self.is_recording = True
                    print(f"开始录制到: {video_path}")
        elif key == ord('1'):
            self.processing_mode = 0
            print("切换到原始视频模式")
        elif key == ord('2'):
            self.processing_mode = 1
            print("切换到灰度效果模式")
        elif key == ord('3'):
            self.processing_mode = 2
            print("切换到边缘检测模式")
        elif key == ord('4'):
            self.processing_mode = 3
            print("切换到模糊效果模式")
    
    def stop_stream(self):
        """停止视频流线程"""
        if not self.running:
            print("视频流未运行")
            return
        
        self.running = False
        
        # 等待线程结束
        if self.stream_thread and self.stream_thread.is_alive():
            self.stream_thread.join(timeout=2.0)
            if self.stream_thread.is_alive():
                print("警告：视频流线程未能正常结束")
        
        # 确保资源被释放
        if self.is_recording and self.video_writer is not None:
            self.video_writer.release()
            self.video_writer = None
            self.is_recording = False
        
        cv2.destroyAllWindows()
        self.receiver.disconnect()
        print("视频流已停止")
    
    def capture(self):
        """捕获当前视频帧"""
        # 如果视频流未运行，启动它
        if not self.running:
            if not self.start_stream_thread(show_window=False):
                return None
        
        # 返回当前帧（如果有）
        with self.frame_lock:
            if self.current_frame is not None:
                return self.current_frame.copy()
        return None
    
    def get_processed_frame(self, processing_mode=None):
        """获取处理后的当前帧
        
        参数:
            processing_mode: 处理模式 (0: 原始, 1: 灰度, 2: 边缘检测, 3: 模糊)
                            如果为None，则使用当前设置的处理模式
        
        返回:
            处理后的帧，如果没有可用帧则返回None
        """
        # 获取当前帧
        frame = self.capture()
        if frame is None:
            return None
        
        # 使用指定的处理模式或当前模式
        mode = processing_mode if processing_mode is not None else self.processing_mode
        
        # 应用处理
        if mode == 0:
            return frame  # 原始
        elif mode == 1:
            return self.apply_grayscale(frame)  # 灰度
        elif mode == 2:
            return self.apply_edge_detection(frame)  # 边缘检测
        elif mode == 3:
            return self.apply_blur(frame)  # 模糊
        else:
            return frame  # 默认返回原始帧
    
    def save_current_frame(self, filename=None, processing_mode=None):
        """保存当前帧到文件
        
        参数:
            filename: 保存的文件名，如果为None则自动生成
            processing_mode: 处理模式，如果为None则使用当前设置的处理模式
        
        返回:
            保存成功返回True，否则返回False
        """
        # 获取处理后的帧
        frame = self.get_processed_frame(processing_mode)
        if frame is None:
            logger.error("没有可用的帧可保存")
            return False
        
        # 如果没有指定文件名，则自动生成
        if filename is None:
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            filename = os.path.join(self.output_dir, f"frame_{timestamp}.jpg")
        
        # 保存帧
        try:
            cv2.imwrite(filename, frame)
            logger.info(f"已保存帧到: {filename}")
            return True, filename
        except Exception as e:
            logger.info(f"保存帧失败: {str(e)}")
            return False, filename
